import numpy as np
import matplotlib.pyplot as plt

def new_extract_shape(wave,figure_show=True):
    window = 250
    var = []
    lens = wave.shape[0]
    for i in range(0, lens - window):
        var.append(np.var(wave[i:i + window]))
    avg = np.mean(var)
    var = np.array(var)
    a = []
    for i in range(len(var) - 1):
        if var[i] < avg < var[i + 1] or var[i] > avg > var[i + 1]:
            a.append(i)

    nums = []
    tmp = None
    for i in range(len(a) // 2):
        if a[2 * i + 1] - a[2 * i] > window: #表示可能直接是预选区间
            tmp = None
            nums.append(a[2 * i])
            nums.append(a[2 * i + 1])

        else:
            if tmp is not None:
                mm = np.mean(wave[tmp:a[2 * i + 1]])
                if wave[tmp] < mm:  #如果选中的tmp比平均值大，那么说明此时tmp在顶点位置，选错了，去除！！
                    nums.append(tmp)
                    nums.append(a[2 * i + 1])
                    tmp = None
                else:
                    tmp = a[2 * i]
            else:
                tmp = a[2 * i]


    res = []
    k1 = []
    k2 = []
    l_p, r_p, l_m, r_m = [], [], [], []
    for i in range(len(nums) // 2):
        l = nums[i * 2]
        r = nums[i * 2 + 1]
        while True:
            if l - 1 >= 0:
                if l - 1 == 0 or var[l - 1] > var[l] < var[l + 1]:
                    break
                else:
                    l -= 1
        while True:
            if r + 1 < len(var):
                if r + 1 == len(var) - 1 or var[r - 1] > var[r] < var[r + 1]:
                    break
                else:
                    r += 1
        k1.append(l)
        k2.append(r)
        r += window
        l -= window // 3
        if l < 0:
            l = 0
        if r >= lens:
            r = lens
        print(l, r)
        sub_carrier_avg = np.mean(wave[l:r])

        top_index = np.where(wave == np.max(wave[l:r]))[0][0]
        ll = rr = top_index #候选位置
        while ll > l:
            if wave[ll] > sub_carrier_avg:
                ll -= 1
            else:
                break
        while rr < r:
            if wave[rr] > sub_carrier_avg:
                rr += 1
            else:
                break
        print(sub_carrier_avg)
        l_median = np.median(wave[l:ll])
        r_median = np.median(wave[rr:r])

        while ll - 1 > l:
            if wave[ll] > l_median :
                ll -= 1
            else:
                if wave[ll - 1] > wave[ll] < wave[ll + 1]:
                    break
                else:
                    ll -= 1
        while rr + 1 < r:
            if wave[rr] > r_median:
                rr += 1
            else:
                if wave[rr - 1] > wave[rr] < wave[rr + 1]:
                    break
                else:
                    rr += 1
        #     ####### 这里显示划分的子波形区域以及子波形中的平均值
        res.append(wave[ll:rr].tolist())
        l_p.append(ll)
        r_p.append(rr)
        l_m.append(l_median)
        r_m.append(r_median)
        # if figure_show:
        #     plt.plot([ll, rr], [wave[ll], wave[rr]], 'b^')
        #     print(l,ll,r,rr,l_median, r_median)
        #     plt.plot([l, ll], [l_median, l_median], 'r-')
        #     plt.plot([rr, r], [r_median, r_median], 'b--')
            # plt.plot(wave[l:r])
            # plt.plot(wave)
            # plt.plot(range(l, r), wave[l:r], 'r.')
            # plt.plot([l, r], [tmp_mean, tmp_mean])
            # plt.show()
    if figure_show:
        # 画方差图
        # plt.xlim(0, 3500)
        plt.figure(figsize=(12, 4))
        plt.subplot(121)
        plt.plot(var)
        plt.plot(k1, var[k1], 'r^', markersize=8)
        plt.plot(k2, var[k2], 'r*', markersize=10)
        plt.legend(['方差', 'L', 'R'], loc='upper right')
        plt.ylabel('方差值', fontsize=14)
        plt.xlabel('数量', fontsize=14)
        plt.xticks(fontsize=14)
        plt.xticks(fontsize=14)
        plt.tight_layout()
        # plt.savefig('figure4-1.eps', dpi=600, format='eps')
        # plt.show()
    if figure_show:
        plt.subplot(122)
        plt.plot(wave)
        # plt.xlim(0, 3800)
        plt.xlabel('样本数', fontsize=14)
        plt.ylabel('PCA值', fontsize=14)
        plt.xticks(fontsize=14)
        plt.xticks(fontsize=14)
        plt.plot(l_p, wave[l_p], 'r^', markersize=8)
        plt.plot(r_p, wave[r_p], 'r*', markersize=10)
        plt.legend(['CSI值', '起始点', '终止点'], loc='upper right')
        plt.tight_layout()
        plt.savefig('figure4-2.eps', dpi=600, format='eps')
        plt.show()
    return res


